{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "39f9cef1",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/aiforsec22/IEEEEuroSP23/blob/main/notebooks/malware-similarity.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ee6cd31",
   "metadata": {
    "id": "AKZVPZlClCjq"
   },
   "source": [
    "### Installing dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e28f9368",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "UZJt0nICk-Hq",
    "outputId": "b4ef8eb2-9664-44cc-afff-063237d988ea"
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/aiforsec/LADDER.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b97f4e5",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "VW0V59IWk-Lw",
    "outputId": "cefab1cd-aae6-40ae-b9ee-f652b8a02c42"
   },
   "outputs": [],
   "source": [
    "%cd LADDER/attack_pattern/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1931a63",
   "metadata": {
    "id": "FPPyqtJAlAf7"
   },
   "source": [
    "### Import modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49d50eb6",
   "metadata": {
    "id": "69db5177"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from sklearn import metrics\n",
    "from sklearn.datasets import make_blobs\n",
    "from sklearn.preprocessing import StandardScaler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66f49e3b",
   "metadata": {
    "id": "2c4ff716"
   },
   "source": [
    "### Read all malware, threat actor and triples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05b0afc2",
   "metadata": {
    "id": "5fc17d4b"
   },
   "outputs": [],
   "source": [
    "with open('all_malware.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()\n",
    "malware = []\n",
    "\n",
    "for line in text.split('\\n')[:-1]:\n",
    "    malware.append(line)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b461bd7c",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "f523b472",
    "outputId": "cbcdc925-128b-465b-dd22-5eb2e23283fb"
   },
   "outputs": [],
   "source": [
    "len(malware)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61b71670",
   "metadata": {
    "id": "7e7fe1b3"
   },
   "outputs": [],
   "source": [
    "with open('all_threat_actors.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()\n",
    "actors = []\n",
    "\n",
    "for line in text.split('\\n')[:-1]:\n",
    "    actors.append(line)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e382487",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "82af03e1",
    "outputId": "95038c23-b8ce-4781-9c37-9712da2b00e4"
   },
   "outputs": [],
   "source": [
    "len(actors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f44f8dc8",
   "metadata": {
    "id": "d009bd20"
   },
   "outputs": [],
   "source": [
    "def read_triples(fname):\n",
    "    triples = []\n",
    "    with open(fname, 'r', encoding='utf-8') as f:\n",
    "        text = f.read()\n",
    "\n",
    "    for line in text.split('\\n'):\n",
    "        if len(line) > 0:\n",
    "            e1, r, e2 = line.split('\\t')\n",
    "            triples.append([e1, r, e2])\n",
    "    return triples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c659cd08",
   "metadata": {
    "id": "3962a0d0"
   },
   "outputs": [],
   "source": [
    "triples = read_triples('150_all.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f990ad85",
   "metadata": {
    "id": "f5fea0d5"
   },
   "outputs": [],
   "source": [
    "def get_malware_nodes(mal, triples):\n",
    "    nodes = set()\n",
    "    \n",
    "    for e1, r, e2 in triples:\n",
    "        if e1 == mal:\n",
    "            nodes.add((e2, r))\n",
    "        elif e2 == mal:\n",
    "            nodes.add((e1, r))\n",
    "    return nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a2f0c9e",
   "metadata": {
    "id": "48f94fdc"
   },
   "outputs": [],
   "source": [
    "def get_all_malware_nodes(triples):\n",
    "    malware_nodes = {}\n",
    "    for m in malware:\n",
    "        nodes = get_malware_nodes(m, triples)\n",
    "        if len(nodes) > 0:\n",
    "            malware_nodes[m] = nodes\n",
    "    return malware_nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56290d0f",
   "metadata": {
    "id": "ac298e88"
   },
   "outputs": [],
   "source": [
    "def get_all_actor_nodes(triples):\n",
    "    actor_nodes = {}\n",
    "    for m in actors:\n",
    "        nodes = get_malware_nodes(m, triples)\n",
    "        nodes_list = list(nodes)\n",
    "        for x in nodes_list:\n",
    "            if x[1] == 'hasAuthor':\n",
    "                mal_nodes = get_malware_nodes(x[0], triples)\n",
    "                for z in mal_nodes:\n",
    "                    if z[1] in ['targets', 'uses', 'exploits', 'indicates', 'isA', 'variantOf',]:\n",
    "                        nodes.add(z)\n",
    "#                     else:\n",
    "#                         nodes.add(x[0], )\n",
    "        if len(nodes) > 0:\n",
    "            actor_nodes[m] = nodes\n",
    "    return actor_nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da5ed482",
   "metadata": {
    "id": "c588e6bd"
   },
   "outputs": [],
   "source": [
    "malware_nodes = get_all_malware_nodes(triples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4f9554a",
   "metadata": {
    "id": "20452b0d"
   },
   "outputs": [],
   "source": [
    "def get_distance(node1, node2, type='jaccard'):\n",
    "    union = node1.union(node2)\n",
    "    intersect = node1.intersection(node2)\n",
    "    if type == 'intersect':\n",
    "        return 1000-len(intersect)\n",
    "    elif type == 'jaccard':\n",
    "        return 1 - len(intersect)/len(union)\n",
    "    elif type == 'overlap':\n",
    "        return 1 - len(intersect)/min(len(node1), len(node2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4e9a7b6",
   "metadata": {
    "id": "1f07c074"
   },
   "outputs": [],
   "source": [
    "def find_most_similar_malware(mal, triple_fname):\n",
    "    triples = read_triples(triple_fname)\n",
    "    malware_nodes = get_all_malware_nodes(triples)\n",
    "\n",
    "    malware_list = list(malware_nodes.keys())\n",
    "    mal_node_i = malware_nodes[mal]\n",
    "        \n",
    "    dist = []    \n",
    "    for j in range(len(malware_list)):\n",
    "            mal_node_j = malware_nodes[malware_list[j]]\n",
    "            dist.append([malware_list[j], get_distance(mal_node_i, mal_node_j, 'jaccard')])\n",
    "    dist.sort(key=lambda x: x[1])\n",
    "    \n",
    "    return dist[1:6]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32d5add7",
   "metadata": {
    "id": "5602bb27"
   },
   "source": [
    "### Find the malware most similar to FluBot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b01159e4",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0187348d",
    "outputId": "78abc781-2964-43ad-c009-eb8aea564029"
   },
   "outputs": [],
   "source": [
    "find_most_similar_malware('FluBot', '12k_all.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b0de8b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# should print the following\n",
    "# [['TeaBot', 0.7906976744186046],\n",
    "#  ['Medusa', 0.8064516129032258],\n",
    "#  ['Gustuff', 0.8115942028985508],\n",
    "#  ['Ghimob', 0.823943661971831],\n",
    "#  ['Faketoken', 0.8260869565217391]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fd8154e",
   "metadata": {
    "id": "7a95eebc"
   },
   "outputs": [],
   "source": [
    "triples = read_triples('12k_all.txt')\n",
    "malware_nodes = get_all_malware_nodes(triples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0a93bca",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "b1e146da",
    "outputId": "21390299-d18f-412d-fd02-6be9f6b0a2e2"
   },
   "outputs": [],
   "source": [
    "mal_node_i = malware_nodes['FluBot']\n",
    "mal_node_j = malware_nodes['TeaBot']\n",
    "\n",
    "print(mal_node_i.intersection(mal_node_j))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cd0d725",
   "metadata": {
    "id": "a17cce9f"
   },
   "outputs": [],
   "source": [
    "def find_most_similar_threat_actor(act, triple_fname):\n",
    "    triples = read_triples(triple_fname)\n",
    "    actor_nodes = get_all_actor_nodes(triples)\n",
    "    \n",
    "    actor_list = list(actor_nodes.keys())\n",
    "    actor_node_i = actor_nodes[act]\n",
    "    \n",
    "    dist = []    \n",
    "    for j in range(len(actor_list)):\n",
    "            actor_node_j = actor_nodes[actor_list[j]]\n",
    "            dist.append([actor_list[j], get_distance(actor_node_i, actor_node_j, 'jaccard')])\n",
    "    dist.sort(key=lambda x: x[1])\n",
    "    \n",
    "    return dist[1:6]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba0c82e3",
   "metadata": {
    "id": "bc9e4643"
   },
   "source": [
    "### Find the most similar threat actor to APT15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e31cda99",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "254302e0",
    "outputId": "e02ea942-9205-43da-a68a-b31ea146c788"
   },
   "outputs": [],
   "source": [
    "find_most_similar_threat_actor('APT15', '12k_all.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e8ef825",
   "metadata": {},
   "outputs": [],
   "source": [
    "# should print the folliwng\n",
    "# [['GREF', 0.5333333333333333],\n",
    "#  ['Boyusec', 0.574468085106383],\n",
    "#  ['Ke3chang', 0.5833333333333333],\n",
    "#  ['APT-C-50', 0.8163265306122449],\n",
    "#  ['Kitten', 0.8333333333333334]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76460414",
   "metadata": {
    "id": "b7d268ce"
   },
   "outputs": [],
   "source": [
    "triples = read_triples('12k_all.txt')\n",
    "actor_nodes = get_all_actor_nodes(triples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6b5d101",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4c676ee9",
    "outputId": "1be7ebf7-e3e7-4353-aeac-2f390990ad10"
   },
   "outputs": [],
   "source": [
    "mal_node_i = actor_nodes['APT15']\n",
    "mal_node_j = actor_nodes['Boyusec']\n",
    "\n",
    "print(mal_node_i.intersection(mal_node_j))"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "include_colab_link": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
